import numpy
from scipy import sparse
from collections import defaultdict
from math import sin, atan as _atan, pi, cos

"""A simple power flow solver."""

('line information of the form:'
 'i     j       Rij     Xij     Bij')

('bus information of the form:'
 'i     type    V   angle   Pgen    Qgen    Pload   Qload')

# bus types can put in a constants file.
SWING, GEN, LOAD = 0, 1, 2

def tupled_args(fn):
    def tupfunc(*args):
        return fn(args)
    return tupfunc

hs = tupled_args(numpy.hstack)
vs = tupled_args(numpy.vstack)

def atan(imag, real):
    r = _atan(imag / real)
    if real < 0 and imag > 0:
        return r + pi
    if real < 0 and imag < 0:
        return r - pi
    return r

def update_businfo(businfo, a, v):
    "update the businfomation matrix with change in angles and voltages."
    for info in businfo:
        if info[1] == SWING:
            continue
        info[3] = a[0]
        a = a[1:]
        if info[1] == LOAD:
            info[2] = v[0]
            v = v[1:]
    return businfo

def solve_powerflow(businfo, lineinfo, tolerance=1e-4, maxiter=100):
    "solve the power flow given initial businfo and lineinfo."
    error = 10
    y = ymatrix(lineinfo)
    while abs(error) > tolerance and maxiter > 0:
        maxiter -= 1
        # calculate the jacobian.
        jac = jacobian(businfo, y)
        # calculate the change in real and reactive power.
        p, q = change_in_pq(businfo, y)
        # the error is the largest absolute change in P or Q.
        error = abs(vs(p, q)).max()
        # calculate angles and voltages for the next iteration.
        a, v = initial_args(jac, vs(p, q), businfo)
        # now update businfo with this information.
        businfo = update_businfo(businfo, a, v)

    # now calculate the power on swing bus and reactive powers.
    for info in businfo:
        if info[1] == LOAD:
            continue # already know P and Q.
        # calculate the reactive power.
        q = calcPorQ(0, info[2], info[3], info[0], businfo, y, sin)
        info[5] = -q # calcPorQ inverts true value.
        if info[1] == SWING:
            # calculate power for swing bus only.
            p = calcPorQ(0, info[2], info[3], info[0], businfo, y, cos)
            info[4] = -p

    return businfo

def calcPorQ(P, V, iangle, ibusno, bus, ymatrix, sc):
    "calculate change in real or reactive power due to some initial conditions."
    calcs = []
    for info in bus:
        jbusno, jbustype, jV, jangle, Pg, Qg, Pl, Ql = info
        y = ymatrix[ibusno, jbusno]
        calcs.append(jV * abs(y) * sc(iangle - jangle -atan(y.imag, y.real)))
    return P - V * sum(calcs)

def change_in_pq(bus, y):
    "Calculate the delta P and Q values for a given initial condition."
    deltaP, deltaQ = [], []
    for info in bus:
        ibusno, ibustype, iV, iangle, ipg, iqg, ipl, iql = info
        if ibustype == SWING:
            continue
        deltaP.append(calcPorQ(ipg - ipl, iV, iangle, ibusno, bus, y, cos))
        # if gen type is PV, they have a known V hence no need to solve Q.
        if ibustype == LOAD:
            # type is PQ, calc both P and Q.
            deltaQ.append(calcPorQ(iqg - iql, iV, iangle, ibusno, bus, y, sin))
    return vs(*deltaP), vs(*deltaQ)

def initial_args(jacobian, deltaPQ, bus):
    """return an array representing the next initial voltage and angle deltas.
        [angle0, angle1, ... anglen, V0, V1, ... VN
       given the jacobian (square) and delta P and Q matrix."""

    angles, voltages = [], []
    for info in bus:
        _, bustype, V, angle, _, _, _, _ = info
        if bustype == SWING:
            continue
        angles.append(angle)
        if bustype == LOAD:
            voltages.append(V)
    init = vs(vs(*angles), vs(voltages))

    solution = init - numpy.linalg.solve(jacobian, deltaPQ)
    return solution[:len(angles)], solution[len(angles):]

def jacobian(bus, y):
    "A function to create the jacobian given bus information and the ymatrix"
    def n(sign, sc):
        "sign to multiply result and sc is sign or cos function."
        def makedelta(iV, jV, ia, ja, y):
            return sign * iV * jV * abs(y) * sc(ia - ja - atan(y.imag, y.real))
        return makedelta

    j1 = JacobianJ1(bus, y, n(-1, sin), n(1, sin), n(-1,sin))

    def myself(iV, jV, ia, ja, y):
        return -1 * jV * abs(y) * cos(ia - ja - atan(y.imag, y.real))

    def other(iV, jV, ia, ja, y):
        return -1 * iV * abs(y) * cos(ia - ja - atan(y.imag, y.real))

    j2 = JacobianJ1(bus, y, other, myself, myself, (GEN,))

    j3 = JacobianJ1(bus, y, n(1, cos), n(-1, cos), n(1, cos), ignorerow=GEN)

    def myself(iV, jV, ia, ja, y):
        return -jV * abs(y) * sin(ia - ja - atan(y.imag, y.real))

    def other(iV, jV, ia, ja, y):
        return -iV * abs(y) * sin(ia - ja - atan(y.imag, y.real))

    j4 = JacobianJ1(bus, y, other, myself, myself, ignorecol=GEN,
            ignorerow=GEN)

    return vs(hs(j1, j2), hs(j3, j4))

def JacobianJ1(businfo, ymatrix, fother, fself, fextra, ignorecol=None,
                                                        ignorerow=None):
    """create part of the jacobian given the bus information and ymatrix.
      fother - a function that accepts Vi and Vj, angle_i and angle_j
               and admittance matrix y and combines them to find the
               partial derivative of change in power or change in
               reactive power for a change in angle or voltage of a
               connected bus.
      fself - a function that accepts Vi and Vj, angle_i and angle_j and
              admittance matrix y and combines them to find the partial
              derivative of change in power or change in reactive power
              for a change in angle or voltage of a bus for which i==j i.e
              a diagonal element.
      fextra - a function with the same signature as fother and fself that
               will calculate the additional component of the partial
               derivative of change in power or reactive power when considering
               changes at bus i==j i.e a diagonal element. such that the
               equation looks like:
               element = sum(fself(for all connected buses)) + fextra(bus i==j)
      ignorecol is a column type (i.e GEN or LOAD) to ignore i.e not
      create a column when the bus j (connected bus) has this type.
      
      ignorerow is a type (i.e GEN or LOAD) where a row will not be created for
      bus i that has this type, i.e the 'from' bus. """
    jacobian = list()

    for info in businfo:
        # for each row.
        ibusno, ibustype, iV, iangle, ipg, iqg, ipl, iql = info

        # ignore if a slackbus, no jacobian elements here.
        if ibustype == SWING or ibustype == ignorerow:
            continue

        jacobian.append([])

        for info in businfo:
            # for each column.
            jbusno, jbustype, jV, jangle, jpg, jqg, jpl, jql = info

            if jbustype == ignorecol or jbustype == SWING:
                continue

            if jbusno != ibusno:
                y = ymatrix[ibusno, jbusno]
                delta = fother(iV, jV, iangle, jangle, y)
                jacobian[-1].append(delta)
                continue

            delta = 0
            for info in businfo:
                # for each cell.
                kbusno, kbustype, kV, kangle, kpg, kqg, kpl, kql = info

                y = ymatrix[ibusno, kbusno]
                if y == 0:
                    continue 
                
                if kbusno == ibusno:
                    delta += fextra(iV, kV, iangle, kangle, y)

                delta += fself(iV, kV, iangle, kangle, y)

            jacobian[-1].append(delta)
    return numpy.array(jacobian)

def ymatrix(lineinfo):
    'form the ymatrix from given line information.'

    # the number of buses in our system.
    size = lineinfo[:, :2].max() + 1

    # create a complex sparse matrix to hold the results.
    a = sparse.lil_matrix((size, size), dtype=complex)

    # to store all the shunt values.
    shunts = numpy.zeros(size, dtype=complex)

    # fill in all the off diagonal components.
    for i, j, R, X, B in lineinfo:
        # assume that lineinfo can be neatly unpacked.
        a[i, j] = a[j, i] = -1 / complex(R, X)

        # save the shunt values for later.
        # split half of the shunt and place at either end of line.
        shunts[i] += complex(0, B) / 2
        shunts[j] += complex(0, B) / 2

    # create a column of ones. It is used to sum each row.
    ones = numpy.ones(size)

    # sum each of the admittances and add the shunt admittances.
    rowsums = -1 * a * ones + shunts

    # diagonals are the sum of all admittances to bus i.
    a.setdiag(rowsums)

    return a
